#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
结果可视化工具
生成论文级别的对比图表

使用方法:
    python visualize_results.py --results evaluation_results/summary_*.txt
"""

import re
import matplotlib.pyplot as plt
import numpy as np
import argparse
from pathlib import Path

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

# 论文配置
plt.rcParams['figure.dpi'] = 300
plt.rcParams['savefig.dpi'] = 300
plt.rcParams['font.size'] = 10


def parse_summary_file(summary_path):
    """解析汇总文件"""
    with open(summary_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    results = {}
    
    # 提取 BLEU 分数
    pattern = r'(\w+(?:\s+\w+)*?)\s+\|\s+([\d.]+|N/A)'
    matches = re.findall(pattern, content)
    
    for model_name, score in matches:
        model_name = model_name.strip()
        if score != "N/A":
            results[model_name] = float(score)
    
    return results


def plot_bleu_comparison(results, output_path="figures/bleu_comparison.png"):
    """绘制 BLEU 分数对比柱状图"""
    models = list(results.keys())
    scores = list(results.values())
    
    fig, ax = plt.subplots(figsize=(10, 6))
    
    # 颜色方案
    colors = ['#E74C3C', '#3498DB', '#2ECC71', '#F39C12', '#9B59B6', '#1ABC9C']
    
    bars = ax.bar(range(len(models)), scores, color=colors[:len(models)], alpha=0.8, edgecolor='black')
    
    # 添加数值标签
    for bar, score in zip(bars, scores):
        height = bar.get_height()
        ax.text(bar.get_x() + bar.get_width()/2., height,
                f'{score:.2f}',
                ha='center', va='bottom', fontsize=9, fontweight='bold')
    
    ax.set_xlabel('Model', fontsize=12, fontweight='bold')
    ax.set_ylabel('BLEU Score', fontsize=12, fontweight='bold')
    ax.set_title('BLEU Score Comparison Across Different Optimization Methods', 
                 fontsize=14, fontweight='bold', pad=20)
    ax.set_xticks(range(len(models)))
    ax.set_xticklabels(models, rotation=15, ha='right')
    ax.set_ylim(0, max(scores) * 1.2)
    ax.grid(axis='y', alpha=0.3, linestyle='--')
    
    plt.tight_layout()
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    plt.savefig(output_path, bbox_inches='tight')
    print(f"[SAVED] {output_path}")
    plt.close()


def plot_improvement_percentage(results, baseline_key="Baseline (原始)", 
                                output_path="figures/improvement_percentage.png"):
    """绘制相对提升百分比"""
    if baseline_key not in results:
        print(f"[SKIP] Baseline '{baseline_key}' not found")
        return
    
    baseline_score = results[baseline_key]
    improvements = {}
    
    for model, score in results.items():
        if model != baseline_key:
            improvement = ((score - baseline_score) / baseline_score) * 100
            improvements[model] = improvement
    
    models = list(improvements.keys())
    percentages = list(improvements.values())
    
    fig, ax = plt.subplots(figsize=(10, 6))
    
    colors = ['#2ECC71' if p > 0 else '#E74C3C' for p in percentages]
    bars = ax.barh(range(len(models)), percentages, color=colors, alpha=0.8, edgecolor='black')
    
    # 添加数值标签
    for bar, pct in zip(bars, percentages):
        width = bar.get_width()
        ax.text(width, bar.get_y() + bar.get_height()/2.,
                f'{pct:+.1f}%',
                ha='left' if pct > 0 else 'right',
                va='center', fontsize=9, fontweight='bold')
    
    ax.set_xlabel('Improvement (%)', fontsize=12, fontweight='bold')
    ax.set_ylabel('Model', fontsize=12, fontweight='bold')
    ax.set_title(f'Performance Improvement Relative to Baseline', 
                 fontsize=14, fontweight='bold', pad=20)
    ax.set_yticks(range(len(models)))
    ax.set_yticklabels(models)
    ax.axvline(0, color='black', linewidth=1, linestyle='-')
    ax.grid(axis='x', alpha=0.3, linestyle='--')
    
    plt.tight_layout()
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    plt.savefig(output_path, bbox_inches='tight')
    print(f"[SAVED] {output_path}")
    plt.close()


def plot_model_size_comparison(output_path="figures/model_size_comparison.png"):
    """绘制模型大小对比"""
    models = ['Baseline\n(60M)', 'Standard\nFine-tune\n(60M)', 
              'Distillation\n(60M)', 'Multi-Task\n(60M)', 
              'Contrastive\n(60M)', 'RL\n(60M)']
    
    sizes = [60, 60, 60, 60, 60, 60]  # 单位：百万参数
    lora_params = [0, 0.3, 0.3, 0.3, 0.3, 0.3]  # LoRA 额外参数
    
    fig, ax = plt.subplots(figsize=(10, 6))
    
    x = np.arange(len(models))
    width = 0.35
    
    bars1 = ax.bar(x - width/2, sizes, width, label='Base Parameters', 
                   color='#3498DB', alpha=0.8, edgecolor='black')
    bars2 = ax.bar(x + width/2, lora_params, width, label='LoRA Parameters', 
                   color='#E67E22', alpha=0.8, edgecolor='black')
    
    ax.set_xlabel('Model', fontsize=12, fontweight='bold')
    ax.set_ylabel('Parameters (Million)', fontsize=12, fontweight='bold')
    ax.set_title('Model Size Comparison', fontsize=14, fontweight='bold', pad=20)
    ax.set_xticks(x)
    ax.set_xticklabels(models, fontsize=9)
    ax.legend(fontsize=10)
    ax.grid(axis='y', alpha=0.3, linestyle='--')
    
    plt.tight_layout()
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    plt.savefig(output_path, bbox_inches='tight')
    print(f"[SAVED] {output_path}")
    plt.close()


def plot_training_loss_curve(log_file="model/text2code_lora/training_log.csv",
                             output_path="figures/training_loss.png"):
    """绘制训练损失曲线"""
    try:
        import pandas as pd
        df = pd.read_csv(log_file)
        
        fig, ax = plt.subplots(figsize=(10, 6))
        
        ax.plot(df['step'], df['loss'], color='#E74C3C', linewidth=2, label='Training Loss')
        
        # 添加平滑曲线
        from scipy.ndimage import uniform_filter1d
        smoothed = uniform_filter1d(df['loss'], size=20)
        ax.plot(df['step'], smoothed, color='#3498DB', linewidth=2, 
                linestyle='--', label='Smoothed Loss')
        
        ax.set_xlabel('Training Steps', fontsize=12, fontweight='bold')
        ax.set_ylabel('Loss', fontsize=12, fontweight='bold')
        ax.set_title('Training Loss Curve', fontsize=14, fontweight='bold', pad=20)
        ax.legend(fontsize=10)
        ax.grid(alpha=0.3, linestyle='--')
        
        plt.tight_layout()
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        plt.savefig(output_path, bbox_inches='tight')
        print(f"[SAVED] {output_path}")
        plt.close()
    
    except Exception as e:
        print(f"[SKIP] 无法绘制训练曲线: {e}")


def generate_latex_table(results, output_path="figures/results_table.tex"):
    """生成 LaTeX 表格"""
    latex_code = r"""
\begin{table}[ht]
\centering
\caption{Performance Comparison of Different Optimization Methods}
\label{tab:results}
\begin{tabular}{lc}
\toprule
\textbf{Method} & \textbf{BLEU Score} \\
\midrule
"""
    
    for model, score in results.items():
        latex_code += f"{model} & {score:.2f} \\\\\n"
    
    latex_code += r"""
\bottomrule
\end{tabular}
\end{table}
"""
    
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(latex_code)
    
    print(f"[SAVED] {output_path}")


def main(args):
    print("=" * 50)
    print("   结果可视化工具")
    print("=" * 50)
    print("")
    
    # 查找最新的汇总文件
    if args.results:
        summary_file = args.results
    else:
        result_files = sorted(Path("evaluation_results").glob("summary_*.txt"))
        if not result_files:
            print("[ERROR] 未找到汇总文件！请先运行 evaluate_all_models.sh")
            return
        summary_file = result_files[-1]
    
    print(f"读取汇总文件: {summary_file}")
    results = parse_summary_file(summary_file)
    
    print(f"找到 {len(results)} 个模型结果:")
    for model, score in results.items():
        print(f"  - {model}: {score:.2f}")
    print("")
    
    # 生成图表
    print("生成可视化图表...")
    
    plot_bleu_comparison(results)
    plot_improvement_percentage(results)
    plot_model_size_comparison()
    plot_training_loss_curve()
    generate_latex_table(results)
    
    print("")
    print("[SUCCESS] 所有图表已生成！")
    print("输出目录: figures/")
    print("")
    print("图表列表:")
    print("  - bleu_comparison.png         (BLEU 分数对比)")
    print("  - improvement_percentage.png  (相对提升百分比)")
    print("  - model_size_comparison.png   (模型大小对比)")
    print("  - training_loss.png           (训练损失曲线)")
    print("  - results_table.tex           (LaTeX 表格)")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--results", type=str, default=None,
                        help="汇总文件路径（默认使用最新的）")
    
    args = parser.parse_args()
    main(args)
